import os
import cv2
import numpy as np
import matplotlib.pyplot as plt

# 定义类别名称
group_id_to_category = {
    0: "box",
    1: "arrows_right",
    2: "arrows_left"
}

def draw_annotations(image_path):
    annotation_path = image_path.replace(".jpg", ".txt")

    if not os.path.exists(annotation_path):
        print(f"Annotation file not found for {image_path}.")
        return

    image = cv2.imread(image_path)
    with open(annotation_path, 'r') as file:
        lines = file.readlines()

        for line in lines:
            values = line.strip().split(' ')
            class_id, center_x, center_y, width, height = map(float, values[0:5])
            keypoints = list(map(float, values[5:]))

            # 反归一化边界框坐标
            image_height, image_width, _ = image.shape
            center_x = int(center_x * image_width)
            center_y = int(center_y * image_height)
            width = int(width * image_width)
            height = int(height * image_height)

            # 在图像上绘制矩形框
            cv2.rectangle(image, (int(center_x - width/2), int(center_y - height/2)), (int(center_x + width/2), int(center_y + height/2)), (0, 255, 0), 2)

            # 在矩形框周围绘制类别名称
            category = group_id_to_category.get(int(class_id), "Unknown")
            cv2.putText(image, category, (int(center_x - width/2), int(center_y - height/2 - 10)), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 2)

            # 在关键点附近标出序号
            for i in range(0, len(keypoints), 2):
                keypoint_x = int(keypoints[i] * image_width)
                keypoint_y = int(keypoints[i+1] * image_height)
                # keypoint_visibility = keypoints[i+2]
                #
                # if keypoint_visibility == 2:  # 如果关键点可见，绘制实心圆
                    # cv2.circle(image, (keypoint_x, keypoint_y), 3, (255, 0, 0), -1)
                # else:  # 如果关键点不可见，绘制空心圆
                cv2.circle(image, (keypoint_x, keypoint_y), 1, (255, 0, 0), -1)
                cv2.putText(image, str(i//2 + 1), (keypoint_x, keypoint_y), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 2)

    # 显示图像
    # cv2.namedWindow(os.path.basename(image_path), 0)
    # cv2.imshow(os.path.basename(image_path), image)
    cv2.namedWindow("output", 0)
    cv2.imshow("output", image)
    # 将 OpenCV 格式的图像转换为 Matplotlib 格式
    # image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    #
    # fig, ax = plt.subplots()
    # ax.imshow(image_rgb)
    # ax.axis('off')
    # plt.show()

image_path = r"/home/champrin/Desktop/record_data/网盘/兵种内录/a已分类/能量机关/24-分区赛-场地道具训练/标注好/a"

# 获取文件夹A中所有图片文件路径
image_files = [os.path.join(image_path, f) for f in os.listdir(image_path) if f.lower().endswith(('.png', '.jpg', '.jpeg', '.gif'))]

# 根据文件名进行排序
sorted_image_files = sorted(image_files, key=lambda x: os.path.basename(x))

# 指定用于验证的图像文件路径
# sorted_image_files = [r"/home/champrin/Desktop/MV-CS016-10UC+DA1041860/red/Video_20250113104416448/yolo_annotations/exchange_red_Video_20250113104416448_27.jpg"]


for file_path in sorted_image_files:
    print(file_path)
    draw_annotations(file_path)
    key = cv2.waitKey(0)

    if key == 27:
        break

# folder_path = "/home/champrin/Downloads/rock/2/yolo_annotations"
#
# # 获取文件夹中所有图片文件的路径
# image_files = [os.path.join(folder_path, file) for file in os.listdir(folder_path) if file.lower().endswith(('.png', '.jpg', '.jpeg'))]
#
# for image_file in image_files:
#     # 读取图片文件
#     image = cv2.imread(image_file)
#
#     if image is not None:
#         draw_annotations(image_file)
#
#         # 等待按键按下
#         key = cv2.waitKey(0)
#         if key == 27:
#             break
#     else:
#         print(f"Failed to read image: {image_file}")
#
# cv2.destroyAllWindows()
# print("All images processed.")